import pygame  # 用于绘制奖励曲线
import turtle  # 用于绘制奖励曲线

def draw_rewards(rewards):
    """绘制奖励曲线"""
    pygame.init()
    screen = pygame.display.set_mode((800, 600))
    pygame.display.set_caption("Reward Curve")
    screen.fill((255, 255, 255))  # 填充白色背景

    # 动态调整 y 轴比例
    max_reward = max(rewards) if rewards else 1  # 防止除零错误
    min_reward = min(rewards) if rewards else 0
    reward_range = max_reward - min_reward if max_reward != min_reward else 1

    # 绘制 x 轴和 y 轴
    pygame.draw.line(screen, (0, 0, 0), (50, 550), (750, 550), 2)  # x 轴
    pygame.draw.line(screen, (0, 0, 0), (50, 50), (50, 550), 2)    # y 轴

    # 绘制 x 轴刻度
    for i in range(0, len(rewards), len(rewards) // 10):  # 每隔 10% 绘制一个刻度
        x = 50 + i * (700 / len(rewards))
        pygame.draw.line(screen, (0, 0, 0), (x, 550), (x, 545), 2)  # 刻度线
        label = f"{i}"
        text_surface = pygame.font.Font(None, 24).render(label, True, (0, 0, 0))
        screen.blit(text_surface, (x - 10, 560))  # 刻度标签

    # 绘制 y 轴刻度
    for i in range(0, 11):  # 绘制 11 个刻度 (0 到 100%)
        y = 550 - i * 50
        pygame.draw.line(screen, (0, 0, 0), (50, y), (55, y), 2)  # 刻度线
        reward_value = min_reward + (i * reward_range / 10)
        label = f"{reward_value:.2f}"
        text_surface = pygame.font.Font(None, 24).render(label, True, (0, 0, 0))
        screen.blit(text_surface, (10, y - 10))  # 刻度标签

    # 计算每个点的位置
    points = []
    for i, reward in enumerate(rewards):
        x = 50 + i * (700 / len(rewards))  # 动态调整 x 轴
        y = 550 - int(((reward - min_reward) / reward_range) * 500)  # 归一化并映射到屏幕高度
        points.append((int(x), int(y)))

    # 绘制平滑的奖励曲线
    pygame.draw.lines(screen, (0, 0, 255), False, points, 2)

    pygame.display.flip()

    # 事件循环，允许用户手动关闭窗口
    running = True
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

    pygame.quit()
